from flask import Flask, request, jsonify
import torch
import numpy as np
from config import Config
from model.towers import TwoTowerModel
from serving.faiss_indexer import FaissIndexer

app = Flask(__name__)
config = Config()
indexer = FaissIndexer(config)

# 加载模型和索引
model = TwoTowerModel(config)
model.load_state_dict(torch.load(config.TOWERS_MODEL_SAVE_PATH))
user_tower = model.user_tower
user_tower.eval()
#user_tower = UserTower(config.NUM_USERS, config.EMBEDDING_DIM)
#user_tower.load_state_dict(torch.load(config.MODEL_SAVE_PATH))
#user_tower.eval()

indexer.load_index()

@app.route('/recommend', methods=['GET'])
def recommend():
    try:
        # 获取请求参数
        user_id = int(request.args.get('user_id'))
        top_k = int(request.args.get('top_k', config.TOP_K))
        
        # 生成用户嵌入
        with torch.no_grad():
            user_emb = user_tower(torch.tensor([user_id])).numpy()
        
        # 召回物品
        item_ids, scores = indexer.search(user_emb, top_k)
        
        # 构建响应
        return jsonify({
            "user_id": user_id,
            "recommendations": [
                {"item_id": int(item_id), "score": float(score)}
                for item_id, score in zip(item_ids[0], scores[0])
            ]
        })
    
    except Exception as e:
        return jsonify({"error": str(e)}), 400

@app.route('/batch_recommend', methods=['POST'])
def batch_recommend():
    try:
        # 获取批量用户ID
        data = request.json
        user_ids = data['user_ids']
        top_k = int(request.args.get('top_k', config.TOP_K))
        
        # 生成用户嵌入
        with torch.no_grad():
            user_emb = user_tower(torch.tensor(user_ids)).numpy()
        
        # 批量召回
        results = []
        for i, user_id in enumerate(user_ids):
            item_ids, scores = indexer.search(user_emb[i:i+1], top_k)
            results.append({
                "user_id": user_id,
                "recommendations": [
                    {"item_id": int(item_id), "score": float(score)}
                    for item_id, score in zip(item_ids[0], scores[0])
                ]
            })
        
        return jsonify({"results": results})
    
    except Exception as e:
        return jsonify({"error": str(e)}), 400

if __name__ == '__main__':
    app.run(host=config.API_HOST, port=config.API_PORT, threaded=True)